from circuit import *
from simplify_tree import SimplifyTree
from exception import AssumptionValidException

class Optimizer():
    def __init__(self, adder_tree) -> None:
        self.__adder_tree = adder_tree
        self.__target_FAs : set[FA] = set()

    def optimize(self) -> None:
        self.__assign_bit_level()

        self.__deal_fas_connect_with_stdcell()
        
        self.__simplify_by_trees()
        
        self.__deal_fas_with_co_must_be_right()

    def __assign_bit_level(self) -> None:
        """
            Assign all the wires' bit level
        """
        for index in self.__adder_tree.addend_bits():
            # Input wires are always right
            wire : Wire = self.__adder_tree.addend_bits()[index]
            # the FA which links to wire's output have the RIGHT input status  
            for circuit in wire.output_circuits:
                if isinstance(circuit, FA):
                    self.__deal_fa_with_input_must_be_right(circuit, wire)

            # start recurion call for wire
            self.__assign_bit_level_from_input(wire, (index % self.__adder_tree.w()))

        for index in self.__adder_tree.sum_bits():
            # Ouput wires should be right
            wire : Wire = self.__adder_tree.sum_bits()[index]
            # If this wire come from a FA. This FA's ouput should be right
            if isinstance(wire.input_circuit, FA):
                self.__deal_fa_with_output_must_be_right(wire.input_circuit, wire, None)

            # start recurion call for wire
            self.__assign_bit_level_from_output(wire, index)

        # Deal the reamined wire of FA which not be assigned bit level
        for fa in self.__adder_tree.FAs():
            level = fa.get_level()
            
            for wire in fa.output_wires.values():
                if wire.bit_level == -1:
                    wire.bit_level = level
                
            for wire in fa.input_wires.values():
                if wire.bit_level == -1:
                    wire.bit_level = level

    def __deal_fas_connect_with_stdcell(self):
        # Loop all STDCELL
        for c in self.__adder_tree.circuits():
            if isinstance(c, FA) == True:
                continue
            
            # For each input wires for c, they must be right
            for w in c.input_wires.values():
                # w must be right
                if isinstance(w.input_circuit, FA):
                    # w come from a FA, so this FA's output must be right
                    self.__deal_fa_with_output_must_be_right(w.input_circuit, w, c)

            # For each output wires for c, they must be right
            for w in c.output_wires.values():
                # w must be right
                for cc in w.output_circuits:
                    # cc is w's output. And if cc is FA, cc's input must be right
                    if isinstance(cc, FA):
                        self.__deal_fa_with_input_must_be_right(cc, w)

    def __simplify_by_trees(self):
        for i in range(self.__adder_tree.e() + self.__adder_tree.w()):
            for fa in self.__target_FAs:
                if i == fa.get_level():
                    tree = SimplifyTree(fa)
                    tree.simplify()

    def __deal_fas_with_co_must_be_right(self):
        for fa in self.__adder_tree.FAs():
            if fa.CO_must_be_right == True and fa.input_status == InputStatus.INVERT:
                fa.CO_opt = True

    def __deal_fa_with_output_must_be_right(self, fa : FA, wire : Wire, from_circuit : Circuit) -> None:
        for c in wire.output_circuits:
            if c is not from_circuit and isinstance(c, FA):
                self.__deal_fa_with_input_must_be_right(c, wire)
        if fa.output_wires['S'] is wire:
            self.__target_FAs.add(wire.input_circuit)
        else:
            fa.CO_must_be_right = True

    def __deal_fa_with_input_must_be_right(self, fa : FA, wire : Wire) -> None:
        fa.input_status = InputStatus.RIGHT
        # for input wire of cc except w
        for ww in fa.input_wires.values():
            if ww is wire:
                continue
            if isinstance(ww.input_circuit, FA):
                self.__deal_fa_with_output_must_be_right(ww.input_circuit, ww, fa)

    def __check_assign_wire(wire : Wire, bit_level) -> None:
        if wire.bit_level == -1:
            wire.bit_level = bit_level
        elif wire.bit_level != bit_level:
            raise AssumptionValidException()

    def __assign_bit_level_from_output(self, wire : Wire, bit_level) -> None:
        Optimizer.__check_assign_wire(wire, bit_level)

        circuit : Circuit = wire.input_circuit
        
        if isinstance(circuit, FA):
            last_bit_level = (bit_level if wire is circuit.output_wires['S'] else (bit_level - 1))
            self.__assign_bit_level_from_output(circuit.input_wires['A'], last_bit_level)
            self.__assign_bit_level_from_output(circuit.input_wires['B'], last_bit_level)
            self.__assign_bit_level_from_output(circuit.input_wires['CI'], last_bit_level)
        
    def __assign_bit_level_from_input(self, wire : Wire, bit_level) -> None:
        Optimizer.__check_assign_wire(wire, bit_level)
    
        for circuit in wire.output_circuits:
            if isinstance(circuit, FA):
                self.__assign_bit_level_from_input(circuit.output_wires['S'], bit_level)
                self.__assign_bit_level_from_input(circuit.output_wires['CO'], bit_level + 1)
